if '__main__' == __name__:
    import torch
    from torch.utils.data import TensorDataset, DataLoader
    import cv2 as cv
    import sys
    import os
    import glob
    import numpy as np
    import matplotlib.pyplot as plt
    from PyCmpltrtok.common import sep
    from PyCmpltrtok.common_np import uint8_to_flt_by_lut
    from PyCmpltrtok.data.cifar10.load_cifar10 import load
    from PyCmpltrtok.common_torch import torch_infer
    from python_ai.category.torch.vgg.vgg16_torch_cifar10_tvts import VGG

    # labels
    labels = load(only_meta=True)
    print(labels)

    # the pic
    sep('Load image')
    # path = '/home/asuspei/large_data/DL1/_many_files/cifar-10_pngs/cifar/test/9*_ship.png'
    # path = '/home/asuspei/large_data/DL1/_many_files/cifar-10_pngs/clipped/*.double.png'
    # path = '/home/asuspei/large_data/DL1/_many_files/cifar-10_pngs/clipped/*.center.png'
    path = '/home/asuspei/large_data/DL1/_many_files/cifar-10_pngs/clipped/*.center.s.png'
    # path = '/home/asuspei/large_data/DL1/_many_files/cifar-10_pngs/clipped/9440_ship.portrait.png'
    CKPT_PATH = '/home/asuspei/my_svn/python_ai/category/torch/vgg/_save/vgg16_torch_cifar10_tvts.py/v3.0/vgg16_torch_cifar10_avg-5-1.pth'
    argc = len(sys.argv)
    if argc >= 2:
        path = sys.argv[1]
    if argc >= 3:
        CKPT_PATH = sys.argv[2]

    def load_image(path):
        img = cv.imread(path, cv.IMREAD_COLOR)
        img = img.transpose([2, 0, 1])[::-1]
        img = uint8_to_flt_by_lut(img, np.float32)
        img = np.ascontiguousarray(img)
        return img

    if '*' not in path:
        img = load_image(path)
        img = img[None]
        nf = 1
        tgt_paths = [path]
    else:
        tgt_paths = glob.glob(path)
        nf = len(tgt_paths)
        print([os.path.split(x)[1] for x in tgt_paths])
        # img = np.zeros([0, 3, 32, 32], dtype=np.float32)
        img = None
        for p in tgt_paths:
            im = load_image(p)[None]
            if img is None:
                img = np.zeros([0, *im.shape[1:]], dtype=np.float32)
            img = np.r_[img, im]
    print(img.shape)
    # img = img[:32]
    # print(img.shape)

    sep('Check img')
    cimg = img.copy()
    cimg *= 255
    cimg = cimg.astype(np.uint8)
    cimg = cimg.transpose([0, 2, 3, 1])
    plt.figure(figsize=[8, 8])
    spn = 0
    spr = 5
    spc = 5
    for i in range(spr * spc):
        spn += 1
        if spn > nf:
            break
        plt.subplot(spr, spc, spn)
        plt.axis('off')
        plt.imshow(cimg[i])
        title = tgt_paths[i]
        title = os.path.split(title)[1]
        plt.title(title)
    print('Check and close the plotting window to continue ...')
    plt.show()

    # select device
    sep('cpu or gpu')
    device_id = 'cuda:0' if torch.cuda.is_available() else 'cpu'
    device = torch.device(device_id)
    print(device)

    # the model
    sep('The model')
    model = VGG(10, (3, 32, 32)).to(device)
    print(f'Loading from {CKPT_PATH}')
    sdict = torch.load(CKPT_PATH)
    model.load_state_dict(sdict)
    model.eval()  # EVERY IMPORTANT !!!
    print('Loaded')

    # infer
    sep('infer')
    # im = torch.from_numpy(img).to(device)
    im = torch.Tensor(img).to(device)
    # y = model(im).detach().cpu().numpy()
    y = model(im).cpu().detach().numpy()

    def show_y(y):
        print(y.shape)
        print(y)
        cls = y.argmax(axis=1)
        print(cls)
        sum = (cls == 8).sum()
        xlen = len(y)
        rate = sum / xlen
        print(sum, xlen, rate)

    show_y(y)

    sep('by dl')
    batch_size = 128
    ds_x = TensorDataset(im)
    dl_x = DataLoader(ds_x, batch_size, drop_last=False)
    for bx, in dl_x:
        y = model(bx).detach().cpu().numpy()
        break
    show_y(y)

    # common infer
    sep('common infer')
    y = torch_infer(img, model, device, batch_size)
    show_y(y)
